package main

import (
	"context"
	"encoding/csv"
	"fmt"
	"github.com/milvus-io/milvus-sdk-go/v2/entity"
	"log"
	"os"

	milvus "github.com/milvus-io/milvus-sdk-go/v2/client"
)

type QueryResult struct {
	ID     int64       `json:"id"`
	Fields interface{} `json:"fields"`
}

// 分页查询函数
func paginateQuery(ctx context.Context, client milvus.Client) ([]QueryResult, error) {
	var allResults []QueryResult
	//opt := milvus.SearchQueryOptionFunc(func(option *milvus.SearchQueryOption) {
	//	option.Limit = 10
	//	option.Offset = 1
	//	option.ConsistencyLevel = entity.ClStrong
	//	option.IgnoreGrowing = false
	//	//option.Offset = (int64(pageSize) - 1) * int64(pageSize)
	//})

	test, err := client.Query(
		context.Background(), // ctx
		"galleryNew",         // CollectionName
		[]string{},           // partitionNames
		"id > 0",             // expr
		[]string{
			"id",
			"embedding",
			"image_type",
			"gallery_image_id",
			"path",
			"ai_model",
			"model_version",
			"prompt",
			"width",
			"height",
			"aspectRatio",
			"created_at",
			"deleted_at",
		}, // outputFields
		milvus.WithSearchQueryConsistencyLevel(entity.ClStrong),
		milvus.WithLimit(10),
	)
	if err != nil {
		return nil, err
	}
	fmt.Println(test)
	// 获取 collection 统计信息
	//stats, err := client.GetCollectionStatistics(ctx, collectionName)
	//if err != nil {
	//	log.Fatalf("failed to get collection statistics: %v", err)
	//}
	//// 获取实体数量
	//rowCount := stats["row_count"]
	//fmt.Printf("Collection '%s' 中的实体数量: %s\n", collectionName, rowCount)

	//offset := int64(0)
	//for offset < count {
	//	// 分页查询
	//	results, err := client.Query(ctx, collectionName, []string{}, nil, pageSize, offset)
	//	if err != nil {
	//		return nil, fmt.Errorf("failed to query data: %v", err)
	//	}
	//
	//	for _, result := range results {
	//		allResults = append(allResults, QueryResult{
	//			ID:     result.ID,
	//			Fields: result.Fields,
	//		})
	//	}
	//	offset += int64(pageSize)
	//}

	return allResults, nil
}

// 保存结果为 CSV 文件
func saveToCSV(data []QueryResult, filename string) error {
	file, err := os.Create(filename)
	if err != nil {
		return fmt.Errorf("failed to create file: %v", err)
	}
	defer file.Close()

	writer := csv.NewWriter(file)
	defer writer.Flush()

	// 写入表头
	header := []string{"ID", "Fields"}
	if err := writer.Write(header); err != nil {
		return fmt.Errorf("failed to write header: %v", err)
	}

	// 写入数据
	for _, result := range data {
		record := []string{
			fmt.Sprintf("%d", result.ID),
			fmt.Sprintf("%v", result.Fields), // 将字段数据转换为字符串
		}
		if err := writer.Write(record); err != nil {
			return fmt.Errorf("failed to write record: %v", err)
		}
	}

	fmt.Printf("数据已保存为 %s\n", filename)
	return nil
}

func main() {
	// 连接到 Milvus
	client, err := milvus.NewGrpcClient(
		context.Background(), // ctx
		//"10.10.22.104:19530", // addr
		"192.168.30.80:31530", // addr
	)
	if err != nil {
		log.Fatalf("failed to connect to Milvus: %v", err)
	}
	defer client.Close()

	// 创建上下文
	ctx := context.Background()

	// 执行分页查询
	results, err := paginateQuery(ctx, client)
	if err != nil {
		log.Fatalf("查询失败: %v", err)
	}
	fmt.Println(results)
	//// 保存查询结果为 CSV 文件
	//err = saveToCSV(results, "milvus_data.csv")
	//if err != nil {
	//	log.Fatalf("保存为 CSV 失败: %v", err)
	//}
}
